package top.mcwebsite.markdown.render

import androidx.compose.foundation.ScrollState
import androidx.compose.foundation.background
import androidx.compose.foundation.border
import androidx.compose.foundation.horizontalScroll
import androidx.compose.foundation.layout.Box
import androidx.compose.foundation.layout.fillMaxSize
import androidx.compose.foundation.layout.padding
import androidx.compose.foundation.rememberScrollState
import androidx.compose.runtime.Composable
import androidx.compose.runtime.mutableStateMapOf
import androidx.compose.runtime.remember
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.graphics.Color
import androidx.compose.ui.layout.Layout
import androidx.compose.ui.platform.LocalDensity
import androidx.compose.ui.text.style.TextAlign
import androidx.compose.ui.unit.Constraints
import androidx.compose.ui.unit.Dp
import org.intellij.markdown.ast.ASTNode
import org.intellij.markdown.ast.getTextInNode
import org.intellij.markdown.flavours.gfm.GFMElementTypes
import org.intellij.markdown.flavours.gfm.GFMTokenTypes
import top.mcwebsite.markdown.model.MarkdownNode
import top.mcwebsite.markdown.theme.LocalMarkdownTheme

data class MarkdownTable(
    val headers: List<Cell>,
    val aligns: List<TableAlign>,
    val rows: List<List<Cell>>
)


enum class TableAlign {
    Left,
    Center,
    Right;

    fun toTextAlign(): TextAlign {
        return when (this) {
            Left -> TextAlign.Left
            Center -> TextAlign.Center
            Right -> TextAlign.End
        }
    }

    fun toAlignment(): Alignment {
        return when (this) {
            Left -> Alignment.CenterStart
            Center -> Alignment.Center
            Right -> Alignment.CenterEnd
        }
    }
}

data class Cell(val node: ASTNode)

object TableComponent : MarkdownComponent {

    @Composable
    override fun render(markdownNode: MarkdownNode) {
        val theme = LocalMarkdownTheme.current
        val table = getTable(markdownNode.node, markdownNode.content)
        val tableTheme = theme.tableTheme
        MarkdownTable(
            modifier = Modifier.padding(theme.blockPadding)
                .border(width = tableTheme.cellBorderWidth / 2, color = Color.Black),
            rowCount = table.rows.size + 1,
            columnCount = table.headers.size
        ) { row, column ->
            Box(
                modifier = Modifier.fillMaxSize()
                    .border(tableTheme.cellBorderWidth / 2, color = tableTheme.cellBorderColor)
                    .padding(tableTheme.cellPadding)
                    .apply {
                        if (row == 0) {
                            this.background(color = tableTheme.headerBackgroundColor)
                        }
                    }
            ) {
                if (row == 0) {
                    val cell = table.headers[column]
                    MarkdownText(
                        content = markdownNode.content,
                        astNode = cell.node,
                        textStyle = tableTheme.headerTextStyle,
                        markdownState = markdownNode.state
                    )
                } else {
                    val cell = table.rows[row - 1][column]
                    MarkdownText(
                        modifier = Modifier.align(table.aligns[column].toAlignment()),
                        content = markdownNode.content,
                        astNode = cell.node,
                        textStyle = theme.text.copy(
                            textAlign = table.aligns[column].toTextAlign()
                        ),
                        markdownState = markdownNode.state
                    )
                }
            }
        }
    }

    private fun getTable(node: ASTNode, content: String): MarkdownTable {
        val headers = mutableListOf<Cell>()
        val aligns = mutableListOf<TableAlign>()
        val rows = mutableListOf<List<Cell>>()

        node.children.forEach { childNode ->
            when (childNode.type) {
                GFMElementTypes.HEADER -> {
                    childNode.children.filter { it.type == GFMTokenTypes.CELL }.map {
                        Cell(it)
                    }.apply { headers.addAll(this) }
                }

                GFMTokenTypes.TABLE_SEPARATOR -> {
                    val alignStr = childNode.getTextInNode(content).toString().removeSurrounding("|")
                    alignStr.split("|").map {
                        val str = it.trim()
                        if (str.startsWith(":") && str.endsWith(":")) {
                            TableAlign.Center
                        } else if (str.startsWith(":")) {
                            TableAlign.Left
                        } else if (str.endsWith(":")) {
                            TableAlign.Right
                        } else {
                            TableAlign.Left
                        }
                    }.apply { aligns.addAll(this) }
                }

                GFMElementTypes.ROW -> {
                    val cells = mutableListOf<Cell>()
                    childNode.children.filter { it.type == GFMTokenTypes.CELL }.forEach { cellNode ->
                        cells.add(Cell(cellNode))
                    }
                    rows.add(cells)
                }
            }
        }
        return MarkdownTable(headers, aligns, rows)
    }
}

// copy from https://github.com/sunny-chung/composable-table/blob/main/composable-table/src/commonMain/kotlin/com/sunnychung/lib/android/composabletable/ux/Table.kt
@Composable
fun MarkdownTable(
    modifier: Modifier = Modifier,
    rowCount: Int,
    columnCount: Int,
    maxCellWidthDp: Dp = Dp.Infinity,
    maxCellHeightDp: Dp = Dp.Infinity,
    horizontalScrollState: ScrollState = rememberScrollState(),
    cellContent: @Composable (rowIndex: Int, columnIndex: Int) -> Unit
) {
    val columnWidths = remember { mutableStateMapOf<Int, Int>() }
    val rowHeights = remember { mutableStateMapOf<Int, Int>() }

    val maxCellWidth = if (listOf(Dp.Infinity, Dp.Unspecified).contains(maxCellWidthDp)) {
        Constraints.Infinity
    } else {
        with(LocalDensity.current) { maxCellWidthDp.toPx() }.toInt()
    }
    val maxCellHeight = if (listOf(Dp.Infinity, Dp.Unspecified).contains(maxCellHeightDp)) {
        Constraints.Infinity
    } else {
        with(LocalDensity.current) { maxCellHeightDp.toPx() }.toInt()
    }

    var accumWidths: MutableList<Int>
    var accumHeights: MutableList<Int>

    Box(modifier = modifier) {
        Box(
            modifier = Modifier
                .then(Modifier.horizontalScroll(horizontalScrollState))
        ) {
            Layout(
                content = {
                    (0 until rowCount).forEach { rowIndex ->
                        (0 until columnCount).forEach { columnIndex ->
                            cellContent(rowIndex, columnIndex)
                        }
                    }
                },
            ) { measurables, constraints ->
                val placeables = measurables.mapIndexed { index, it ->
                    val columnIndex = index % columnCount
                    val rowIndex = index / columnCount
                    it.measure(
                        Constraints(
                            minWidth = columnWidths[columnIndex] ?: 0,
                            minHeight = rowHeights[rowIndex] ?: 0,
                            maxWidth = maxCellWidth,
                            maxHeight = maxCellHeight
                        )
                    )
                }

                placeables.forEachIndexed { index, placeable ->
                    val columnIndex = index % columnCount
                    val rowIndex = index / columnCount

                    val existingWidth = columnWidths[columnIndex] ?: 0
                    val maxWidth = maxOf(existingWidth, placeable.width)
                    if (maxWidth > existingWidth) {
                        columnWidths[columnIndex] = maxWidth
                    }

                    val existingHeight = rowHeights[rowIndex] ?: 0
                    val maxHeight = maxOf(existingHeight, placeable.height)
                    if (maxHeight > existingHeight) {
                        rowHeights[rowIndex] = maxHeight
                    }
                }

                accumWidths = mutableListOf(0).apply {
                    (1..columnWidths.size).forEach { i ->
                        this += this.last() + columnWidths[i - 1]!!
                    }
                }
                accumHeights = mutableListOf(0).apply {
                    (1..rowHeights.size).forEach { i ->
                        this += this.last() + rowHeights[i - 1]!!
                    }
                }

                val totalWidth = accumWidths.last()
                val totalHeight = accumHeights.last()

                layout(width = totalWidth, height = totalHeight) {
                    placeables.forEachIndexed { index, placeable ->
                        val columnIndex = index % columnCount
                        val rowIndex = index / columnCount

                        placeable.placeRelative(accumWidths[columnIndex], accumHeights[rowIndex])
                    }
                }
            }
        }
        HorizontalScrollBar(horizontalScrollState, modifier = Modifier.align(Alignment.BottomStart))
    }
}

@Composable
expect fun HorizontalScrollBar(scrollState: ScrollState, modifier: Modifier = Modifier)